#!/usr/bin/env python3
import subprocess
from pathlib import Path
from io import BytesIO
import zstd
from fire import Fire
from queue import SimpleQueue
import pickle
import traceback
from glob import glob
from collections import defaultdict, Counter
from os.path import dirname, join, basename, exists
from datetime import date
import zd
import sys
from multiprocessing import Pool
import gzip
from acora import AcoraBuilder
from dataclasses import dataclass, astuple

import multiprocessing

#
# def error(msg, *args):
#   return multiprocessing.get_logger().error(msg, *args)
#
#
# class LogExceptions(object):
#   def __init__(self, callable):
#       self.__callable = callable
#
#   def __call__(self, *args, **kwargs):
#       try:
#           result = self.__callable(*args, **kwargs)
#       except Exception as e:
#           # Here we add some debugging help. If multiprocessing's
#           # debugging is on, it will arrange to log the traceback
#           error(traceback.format_exc())
#           # Re-raise the original exception so the Pool worker can
#           # clean up
#           raise
#
#       # It was fine, give a normal answer
#       return result
#
# class LoggingPool(Pool):
#     def apply_async(self, func, args=(), kwds={}, callback=None):
#         return Pool.apply_async(self, LogExceptions(func), args, kwds, callback)


@dataclass
class User:
  country: bytes = b''
  lang: bytes = b''
  phone: bytes = b''
  mcc: bytes = b''
  time: bytes = b''
  lang: bytes = b''


USER = User()

utf8 = 'utf8'

KEY_LI = tuple(
  b'"%s":"' % i for i in (
    b'sessionTimestamp',
    b'firstSessionTimestamp',
    b'appVersion',
    b'AndroidId',
    b'userId',
    b'carrier',
    b'deviceModel',
    b'deviceSubModel',
    b'mcc:',
    b'countryISO',
    b'country\\u0026language',  # lang
    b'lang\\u0026mcc',  # pop this
  )
)

BackupRestoreActivity = b"BackupRestoreActivity"

KEY_POS = dict((i, pos) for pos, i in enumerate(KEY_LI))
BUILDER = AcoraBuilder(BackupRestoreActivity, *KEY_LI)

AC = BUILDER.build()

# def zcat(infile):
#   proc = subprocess.Popen(
#     ['pigz', "-dc", infile], stdout=subprocess.PIPE, encoding=None
#   )
#   return proc.stdout


def txt2li(txt):
  r = [b''] * len(KEY_LI)
  active = 0
  for key, pos in AC.findall(txt):
    if key == BackupRestoreActivity:
      active = 1
      continue
    pos = len(key) + pos
    val = txt[pos:txt.find(b'"', pos)]
    if val:
      r[KEY_POS[key]] = val.replace(b',', b';')

  r[0] = r[0][:-3]
  r[1] = r[1][:-3]

  t = r[10].replace(b"\\u0026", b' ').split(b' ', 1)

  if len(t) == 2:
    country, lang = t
    if country:
      r[9] = country
    r[10] = lang
  else:
    r[10] = b''
  t = r[11].replace(b"\\u0026", b' ').split(b' ', 1)
  if len(t) == 2:
    lang, mcc = t
    if mcc:
      r[8] = mcc
    if lang:
      r[10] = lang
  r.pop(11)
  r.append(active)
  return r


def day_count_by_gz(gzfile):
  try:
    return _day_count_by_gz(gzfile)
  except Exception:
    return gzfile + "\t" + traceback.format_exc()


def _day_count_by_gz(gzfile):
  outpath = "/root/pack/" + gzfile.split("/", 3).pop()[:-8]
  if exists(outpath):
    print('cached ', outpath)
    try:
      dctx = zstd.ZstdDecompressor()
      io = BytesIO()
      with open(outpath, 'rb') as ifh:
        dctx.copy_stream(ifh, io)
      r = pickle.loads(io.getvalue())
      if r:
        return r
    except Exception:
      print(gzfile)
      traceback.print_exc()

  key2date = defaultdict(lambda: Counter())
  android_dict = {}
  with gzip.open(gzfile) as gzstream:
    for pos, txt in enumerate(gzstream):
      stop = False
      for i in (b'"Register', b'"VerifySms"'):
        if i in txt:
          stop = True
          break
      if stop:
        continue
      time, create_time, app_version, id, phone, carrier, model, sub_model, mcc, country, lang, active = txt2li(
        txt
      )
      if model == sub_model:
        sub_model = b''
      # print(
      #   time, create_time, app_version, id, phone, carrier, model,
      #   sub_model, mcc, country, lang
      # )
      user = android_dict.get(id)
      if not user:
        user = android_dict[id] = User()
      if lang and not user.lang:
        user.lang = lang
      if country and not user.country:
        user.country = country

      if not phone:
        phone = user.phone

      if active:
        if user.time:
          if user.time > time:
            user.time = time
        else:
          user.time = time
      key = [
        id, phone, create_time, app_version, carrier, model,
        sub_model, mcc
      ]
      key = b",".join(key)
      key2date[key][int(time or 0) // 86400] += 1

      if phone:
        if not user.phone:
          key2 = [
            id, b'', create_time, app_version, carrier, model,
            sub_model, mcc, country, lang
          ]
          key2 = b",".join(key2)
          user.phone = phone
          key2date[key] += key2date[key2]
          del key2date[key2]

      if mcc:
        if not user.mcc:
          key2 = [
            id, phone, create_time, app_version, carrier, model,
            sub_model, b'', country, lang
          ]
          key2 = b",".join(key2)
          user.mcc = mcc
          key2date[key] += key2date[key2]
          del key2date[key2]

      # if pos > 100000:
      #   break
  for uid, u in android_dict.items():
    android_dict[uid] = astuple(u)
  r = (tuple(key2date.items()), android_dict)
  Path(dirname(outpath)).mkdir(parents=True, exist_ok=True)
  with zd.open(outpath, "w") as out:
    out.write(pickle.dumps(r))

  return r


day_set = set()


@Fire
def main(dirpath):

  day_id = defaultdict(lambda: Counter())
  android_dict = {}

  def callback(args):
    if isinstance(args, str):
      print(args)
      raise Exception(dirpath)
    key2date, _dict = args
    for id, user in _dict.items():
      user = User(*user)
      if id in android_dict:
        _user = android_dict[id]
        if _user.time:
          if _user.time > user.time:
            _user.time = user.time
        else:
          _user.time = user.time
        if not _user.country:
          _user.country = user.country
        if not _user.lang:
          _user.lang = user.lang
        if not _user.phone:
          _user.phone = user.phone
      else:
        android_dict[id] = user

    global day_set
    for key, c in key2date:
      day_set |= set(c.keys())
      day_id[key] += c


  # pool = Pool(8)
  #
  # for i in glob(join(dirpath, "*.json.gz")):
  #   pool.apply_async(day_count_by_gz, args=(i, ), callback=callback)
  # pool.close()
  # pool.join()
  # pool = Pool(8)

  for i in glob(join(dirpath, "*.json.gz")):
    callback(day_count_by_gz(i))
  # pool.apply_async(day_count_by_gz, args=(i, ), callback=callback)
  # pool.close()
  # pool.join()

  day_li = tuple(sorted(day_set))
  outfile = "/root/zst/" + basename(dirpath.rstrip("/")) + ".csv.zst"
  print(outfile)
  with zd.open(outfile, "w", level=19) as out:
    out.write(
      b"android_id,phone,create_time,app_version,carrier,model,sub_model,mcc,country,lang,restore_time,"
      + b",".join(
        str(date.fromtimestamp(i * 86400)).encode(utf8)
        for i in day_li
      ) + b"\n"
    )
    for key, c in day_id.items():
      user = android_dict.get(key.split(b",", 1)[0], USER)
      out.write(
        key + b"," + b",".join(
          [
            user.country,
            user.lang,
            user.time,
          ] + list(str(c.get(i, '')).encode(utf8) for i in day_li)
        ) + b"\n"
      )
